'''
Created on 2010-8-18

@author: madlee
'''

from wx.glcanvas import GLCanvas
from kuai.kan.docview import *
    
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

import kuai.kan.kanext as kanext 
from kuai.mol import Bond

try:
    import settings
except:
    from kuai.kan import settings



class Mol3DDocument(Document):
    def __init__(self, doctype, filename):
        Document.__init__(self, doctype, filename)
        
    def load(self):
        ct = []
        colors = []
        coords = []
        with open(self.filename) as file:
            for line in file:
                try:
                    natoms = int(line.strip())
                    break
                except:
                    continue
            ct.append(natoms)
            ct.append(0)
            for i in range(natoms):
                line = file.next()
                tokens = line.split()
                # atomI = glyph.AtomGlyphP(self.glyph)
                number = int(tokens[2])
                ct.append(number)
                coords.append(float(tokens[6]))
                coords.append(float(tokens[7]))
                coords.append(float(tokens[8]))
                
                colors.append(settings.DEFAULT_ELEMENT_COLOR[number][0])
                colors.append(settings.DEFAULT_ELEMENT_COLOR[number][1])
                colors.append(settings.DEFAULT_ELEMENT_COLOR[number][2])
                
            line = file.next()
            nbonds = int(line.strip())
            ct[1] = nbonds
            for i in range(nbonds):
                line = file.next()
                tokens = line.split()
                order =  int(tokens[2])
                if order == 1:
                    ct.append(Bond.SINGLE_BOND)
                elif order == 2:
                    ct.append(Bond.DOUBLE_BOND)
                elif order == 3:
                    ct.append(Bond.TRIPLE_BOND)
                elif order == -2:
                    ct.append(Bond.PARTIAL_BOND)
                else:
                    ct.append(Bond.UNKNOWN_BOND)
                ct.append(int(tokens[0])-1)
                ct.append(int(tokens[1])-1)

        self.ct = kanext.makeConnectTable(ct)
        self.coords = kanext.float_array(coords)
        self.colors = kanext.byte_array(colors)
        self.atom_radius = kanext.float_array([0.3] * self.ct[0])
        _, _, self.center, _, self.radius = kanext.mol_region(self.coords)
        if self.radius < 5:
            self.radius = 5
        
class Mol3DView(View, GLCanvas):
    def __init__(self, doc, *args, **kwargs):
        GLCanvas.__init__(self, *args, **kwargs)
        View.__init__(self, doc)
        
        self.init = False
        # initial mouse position
        self.lastx = self.x = 30
        self.lasty = self.y = 30
        self.size = None
        self.Bind(wx.EVT_ERASE_BACKGROUND, self.OnEraseBackground)
        self.Bind(wx.EVT_SIZE, self.OnSize)
        self.Bind(wx.EVT_PAINT, self.OnPaint)
        self.Bind(wx.EVT_LEFT_DOWN, self.OnMouseDown)
        self.Bind(wx.EVT_LEFT_UP, self.OnMouseUp)
        self.Bind(wx.EVT_MOTION, self.OnMouseMotion)
        
        self.axis = kanext.float_array([1, 0, 0, 0, 1, 0, 0, 0, 1])


    def OnEraseBackground(self, event):
        pass # Do nothing, to avoid flashing on MSW.


    def OnSize(self, event):
        size = self.size = self.GetClientSize()
        self.Refresh(False)
        if self.GetContext():
            self.SetCurrent()
            glViewport(0, 0, size.width, size.height)
            glMatrixMode(GL_PROJECTION)
            glLoadIdentity()
            
            gluPerspective(45.0, float(size.width)/size.height, 0.1, 1000.0);

            glMatrixMode(GL_MODELVIEW)
            glLoadIdentity();                    
        event.Skip()


    def OnPaint(self, event):
        _ = wx.PaintDC(self)
        self.SetCurrent()
        if not self.init:
            self.InitGL()
            self.init = True
        self.OnDraw()


    def OnMouseDown(self, evt):
        self.CaptureMouse()
        self.x, self.y = self.lastx, self.lasty = evt.GetPosition()


    def OnMouseUp(self, evt):
        self.ReleaseMouse()


    def OnMouseMotion(self, evt):
        if evt.Dragging() and evt.LeftIsDown():
            self.lastx, self.lasty = self.x, self.y
            self.x, self.y = evt.GetPosition()
            deltaX, deltaY = self.x - self.lastx, self.y - self.lasty
            if deltaX != 0 or deltaY != 0:
                kanext.rotateAxis(self.axis, float(deltaX)/self.size[1], float(deltaY)/self.size[1]) 
                self.Refresh(False)
            
    def InitGL(self):
        glShadeModel(GL_SMOOTH)
        glClearDepth(1.0)
        glEnable(GL_DEPTH_TEST)
        glDepthFunc(GL_LEQUAL)
        
        glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST)
        glLineWidth(settings.DEFAULT_LINE_WIDTH)
        glClearColor(*settings.DEFAULT_BACK_GROUND)
        
        glLightfv(GL_LIGHT0,  GL_AMBIENT, [0.5, 0.5, 0.5, 1.0])
        glLightfv(GL_LIGHT0,  GL_DIFFUSE, [0.5, 0.5, 0.5, 1.0])
        glLightfv(GL_LIGHT0,  GL_POSITION, [0, 0, -4, 1.0])
        
        glEnable(GL_LIGHTING)
        glEnable(GL_LIGHT0)


    def OnDraw(self):
        # clear color and depth buffers
        self.SetCurrent()
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        glLoadIdentity()
        
        glTranslatef(0, 0, -self.document.radius*2)
        kanext.setAxis(self.axis)
        glTranslatef(-self.document.center.x, -self.document.center.y, -self.document.center.z)

        kanext.drawMol3DLine2(self.document.ct, self.document.coords, self.document.colors)
        # kanext.drawAtomBalls(self.document.ct[0], self.document.coords, self.document.colors, self.document.atom_radius)

        self.SwapBuffers()
        
if __name__ == '__main__':
    app = wx.PySimpleApp(redirect=False)
    frame = wx.Frame(parent=None, id=wx.ID_ANY, size=(800, 600))
    filename = r'D:\Madlee\Job\Direct Force Field\7\TestDPFs\Amber95_DPF\Protein\hie.msd'
    # filename = r'D:\Madlee\Job\Direct Force Field\7\TestDPFs\Amber95_DPF\pdb1atp.msd'
    # filename = r'D:\Madlee\Job\Direct Force Field\7\TestDPFs\Amber95_DPF\large_pdb1.msd'
    # filename = r'C:\Users\madlee\Documents\Job\Direct Force Field\7\TestDPFs\Amber95_DPF\large_pdb1.msd'
    # filename = r'C:\Users\madlee\Documents\Job\Direct Force Field\7\TestDPFs\Amber95_DPF\Protein\hie.msd'
    
    doc = Mol3DDocument(None, filename)
    view = Mol3DView(doc, frame)
    frame.Show()
    app.MainLoop()
